import os
import torch
import torchvision.models as models
from reshape import reshape_model


def export(model_path, output_dir):
    # set the device

    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    print('running on device ' + str(device))

    # load the model checkpoint
    print('loading checkpoint:  ' + model_path)
    checkpoint = torch.load(model_path)

    arch = checkpoint['arch']

    # create the model architecture
    print('using model:  ' + arch)
    model = models.__dict__[arch](pretrained=True)

    # reshape the model's output
    model = reshape_model(model, arch, checkpoint['num_classes'])

    # load the model weights
    model.load_state_dict(checkpoint['state_dict'])

    # add softmax layer

    print('adding nn.Softmax layer to model...')
    model = torch.nn.Sequential(model, torch.nn.Softmax(1))

    model.to(device)
    model.eval()

    print(model)

    # create example image data
    resolution = checkpoint['resolution']
    input = torch.ones((1, 3, resolution, resolution)).cuda()
    print('input size:  {:d}x{:d}'.format(resolution, resolution))

    # format output model path

    output = arch + '.onnx'

    if output_dir:
        output = os.path.join(output_dir, output)

    # export the model
    input_names = ["input_0"]
    output_names = ["output_0"]

    print('exporting model to ONNX...')
    torch.onnx.export(model, input, output, verbose=True, input_names=input_names, output_names=output_names)
    print('model exported to:  {:s}'.format(output))

    return True